"""
行情数据管理器
"""
import asyncio
import json
import threading
from typing import Dict, List, Optional, Callable, Any
from datetime import datetime, timedelta
import pandas as pd
from loguru import logger
import websocket

from .okx_client import OKXClient


class MarketDataManager:
    """行情数据管理器"""
    
    def __init__(self, client: OKXClient, cache_manager: Optional[Any] = None):
        """
        初始化行情数据管理器
        
        Args:
            client: OKX客户端实例
            cache_manager: 缓存管理器（可选）
        """
        self.client = client
        self.cache_manager = cache_manager
        
        # WebSocket相关
        self.ws_url = "wss://wspap.okx.com:8443/ws/v5/public" if client.config.get('sandbox', True) \
                     else "wss://ws.okx.com:8443/ws/v5/public"
        self.ws: Optional[websocket.WebSocketApp] = None
        self.ws_thread: Optional[threading.Thread] = None
        self.ws_connected = False
        self.subscriptions: Dict[str, List[Callable]] = {}  # channel -> callbacks
        
        # 数据缓存
        self.latest_prices: Dict[str, float] = {}  # symbol -> price
        self.kline_cache: Dict[str, pd.DataFrame] = {}  # cache_key -> dataframe
        
        # 心跳
        self.last_heartbeat = datetime.now()
        self.heartbeat_interval = 20  # 秒
        
    def get_klines(self, symbol: str, timeframe: str, 
                   start: Optional[datetime] = None, 
                   end: Optional[datetime] = None,
                   limit: int = 500) -> pd.DataFrame:
        """
        获取K线数据
        
        Args:
            symbol: 交易对，如 'BTC/USDT:USDT'
            timeframe: 时间周期，如 '1m', '5m', '1h', '1d'
            start: 开始时间
            end: 结束时间
            limit: 最大K线数量
            
        Returns:
            K线数据DataFrame，包含列: timestamp, open, high, low, close, volume
        """
        try:
            # 生成缓存键
            cache_key = f"{symbol}_{timeframe}_{start}_{end}_{limit}"
            
            # 检查缓存
            if self.cache_manager:
                cached_data = self.cache_manager.get(cache_key)
                if cached_data is not None:
                    logger.debug(f"从缓存获取K线数据: {symbol} {timeframe}")
                    return cached_data
            
            # 从交易所获取数据
            logger.info(f"获取K线数据: {symbol} {timeframe} limit={limit}")
            
            params = {}
            if start:
                params['since'] = int(start.timestamp() * 1000)
            if end:
                params['until'] = int(end.timestamp() * 1000)
            
            ohlcv = self.client._execute_with_retry(
                self.client.exchange.fetch_ohlcv,
                symbol=symbol,
                timeframe=timeframe,
                limit=limit,
                params=params
            )
            
            # 转换为DataFrame
            df = pd.DataFrame(
                ohlcv,
                columns=['timestamp', 'open', 'high', 'low', 'close', 'volume']
            )
            
            # 转换时间戳
            df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
            df.set_index('timestamp', inplace=True)
            
            logger.info(f"获取K线数据成功: {len(df)} 条")
            
            # 缓存数据
            if self.cache_manager:
                self.cache_manager.set(cache_key, df, expire=60)  # 缓存60秒
            
            return df
            
        except Exception as e:
            logger.error(f"获取K线数据失败: {e}")
            raise
    
    def get_latest_price(self, symbol: str) -> float:
        """
        获取最新价格
        
        Args:
            symbol: 交易对
            
        Returns:
            最新价格
        """
        # 优先从WebSocket缓存获取
        if symbol in self.latest_prices:
            return self.latest_prices[symbol]
        
        # 从交易所获取
        try:
            ticker = self.client._execute_with_retry(
                self.client.exchange.fetch_ticker,
                symbol
            )
            price = float(ticker['last'])
            self.latest_prices[symbol] = price
            return price
        except Exception as e:
            logger.error(f"获取最新价格失败: {e}")
            raise
    
    def subscribe_ticker(self, symbol: str, callback: Callable[[Dict], None]):
        """
        订阅实时行情
        
        Args:
            symbol: 交易对
            callback: 回调函数，接收ticker数据
        """
        # 转换symbol格式 (BTC/USDT:USDT -> BTC-USDT-SWAP)
        inst_id = self._convert_symbol_to_inst_id(symbol)
        channel = f"tickers.{inst_id}"
        
        if channel not in self.subscriptions:
            self.subscriptions[channel] = []
        
        self.subscriptions[channel].append(callback)
        logger.info(f"订阅行情: {symbol} ({inst_id})")
        
        # 启动WebSocket连接
        if not self.ws_connected:
            self._start_websocket()
        else:
            # 发送订阅消息
            self._send_subscribe(channel, inst_id)
    
    def subscribe_kline(self, symbol: str, timeframe: str, 
                       callback: Callable[[Dict], None]):
        """
        订阅K线数据
        
        Args:
            symbol: 交易对
            timeframe: 时间周期
            callback: 回调函数，接收K线数据
        """
        # 转换symbol和timeframe格式
        inst_id = self._convert_symbol_to_inst_id(symbol)
        okx_timeframe = self._convert_timeframe(timeframe)
        channel = f"candle{okx_timeframe}.{inst_id}"
        
        if channel not in self.subscriptions:
            self.subscriptions[channel] = []
        
        self.subscriptions[channel].append(callback)
        logger.info(f"订阅K线: {symbol} {timeframe} ({inst_id})")
        
        # 启动WebSocket连接
        if not self.ws_connected:
            self._start_websocket()
        else:
            # 发送订阅消息
            self._send_subscribe(channel, inst_id)
    
    def unsubscribe(self, symbol: str):
        """
        取消订阅
        
        Args:
            symbol: 交易对
        """
        inst_id = self._convert_symbol_to_inst_id(symbol)
        
        # 移除所有相关订阅
        channels_to_remove = []
        for channel in self.subscriptions.keys():
            if inst_id in channel:
                channels_to_remove.append(channel)
        
        for channel in channels_to_remove:
            del self.subscriptions[channel]
            logger.info(f"取消订阅: {channel}")
        
        # 如果没有订阅了，关闭WebSocket
        if not self.subscriptions and self.ws_connected:
            self._stop_websocket()
    
    def _convert_symbol_to_inst_id(self, symbol: str) -> str:
        """
        转换交易对格式
        
        Args:
            symbol: CCXT格式 'BTC/USDT:USDT'
            
        Returns:
            OKX格式 'BTC-USDT-SWAP'
        """
        # BTC/USDT:USDT -> BTC-USDT-SWAP
        parts = symbol.replace('/', '-').replace(':', '-').split('-')
        if len(parts) >= 2:
            return f"{parts[0]}-{parts[1]}-SWAP"
        return symbol
    
    def _convert_timeframe(self, timeframe: str) -> str:
        """
        转换时间周期格式
        
        Args:
            timeframe: CCXT格式 '1m', '5m', '1h'
            
        Returns:
            OKX格式 '1m', '5m', '1H'
        """
        # OKX使用大写H表示小时
        mapping = {
            '1m': '1m',
            '3m': '3m',
            '5m': '5m',
            '15m': '15m',
            '30m': '30m',
            '1h': '1H',
            '2h': '2H',
            '4h': '4H',
            '6h': '6H',
            '12h': '12H',
            '1d': '1D',
            '1w': '1W',
        }
        return mapping.get(timeframe, timeframe)
    
    def _start_websocket(self):
        """启动WebSocket连接"""
        if self.ws_connected:
            return
        
        logger.info(f"启动WebSocket连接: {self.ws_url}")
        
        self.ws = websocket.WebSocketApp(
            self.ws_url,
            on_open=self._on_ws_open,
            on_message=self._on_ws_message,
            on_error=self._on_ws_error,
            on_close=self._on_ws_close
        )
        
        # 在新线程中运行
        self.ws_thread = threading.Thread(target=self.ws.run_forever, daemon=True)
        self.ws_thread.start()
    
    def _stop_websocket(self):
        """停止WebSocket连接"""
        if self.ws:
            self.ws.close()
            self.ws = None
            self.ws_connected = False
            logger.info("WebSocket连接已关闭")
    
    def _on_ws_open(self, ws):
        """WebSocket连接打开"""
        self.ws_connected = True
        self.last_heartbeat = datetime.now()
        logger.info("WebSocket连接已建立")
        
        # 订阅所有频道
        for channel in self.subscriptions.keys():
            # 从channel提取inst_id
            parts = channel.split('.')
            if len(parts) == 2:
                inst_id = parts[1]
                self._send_subscribe(channel, inst_id)
    
    def _on_ws_message(self, ws, message):
        """处理WebSocket消息"""
        try:
            data = json.loads(message)
            
            # 处理心跳
            if 'event' in data:
                if data['event'] == 'subscribe':
                    logger.debug(f"订阅成功: {data.get('arg', {})}")
                elif data['event'] == 'error':
                    logger.error(f"WebSocket错误: {data}")
                return
            
            # 处理数据推送
            if 'data' in data and 'arg' in data:
                arg = data['arg']
                channel = arg.get('channel', '')
                inst_id = arg.get('instId', '')
                
                # 构建完整channel key
                channel_key = f"{channel}.{inst_id}"
                
                # 调用回调函数
                if channel_key in self.subscriptions:
                    for callback in self.subscriptions[channel_key]:
                        try:
                            callback(data['data'][0])
                        except Exception as e:
                            logger.error(f"回调函数执行失败: {e}")
                
                # 更新价格缓存
                if channel == 'tickers':
                    for item in data['data']:
                        price = float(item.get('last', 0))
                        if price > 0:
                            # 转换inst_id回symbol格式
                            symbol = self._convert_inst_id_to_symbol(inst_id)
                            self.latest_prices[symbol] = price
            
            # 更新心跳时间
            self.last_heartbeat = datetime.now()
            
        except Exception as e:
            logger.error(f"处理WebSocket消息失败: {e}")
    
    def _on_ws_error(self, ws, error):
        """WebSocket错误"""
        logger.error(f"WebSocket错误: {error}")
    
    def _on_ws_close(self, ws, close_status_code, close_msg):
        """WebSocket连接关闭"""
        self.ws_connected = False
        logger.warning(f"WebSocket连接关闭: {close_status_code} {close_msg}")
        
        # 如果还有订阅，尝试重连
        if self.subscriptions:
            logger.info("5秒后尝试重连...")
            threading.Timer(5.0, self._start_websocket).start()
    
    def _send_subscribe(self, channel: str, inst_id: str):
        """
        发送订阅消息
        
        Args:
            channel: 频道名称
            inst_id: 交易对ID
        """
        if not self.ws or not self.ws_connected:
            return
        
        # 提取channel类型
        channel_type = channel.split('.')[0]
        
        subscribe_msg = {
            "op": "subscribe",
            "args": [{
                "channel": channel_type,
                "instId": inst_id
            }]
        }
        
        try:
            self.ws.send(json.dumps(subscribe_msg))
            logger.debug(f"发送订阅消息: {subscribe_msg}")
        except Exception as e:
            logger.error(f"发送订阅消息失败: {e}")
    
    def _convert_inst_id_to_symbol(self, inst_id: str) -> str:
        """
        转换交易对格式
        
        Args:
            inst_id: OKX格式 'BTC-USDT-SWAP'
            
        Returns:
            CCXT格式 'BTC/USDT:USDT'
        """
        # BTC-USDT-SWAP -> BTC/USDT:USDT
        parts = inst_id.split('-')
        if len(parts) >= 2:
            return f"{parts[0]}/{parts[1]}:{parts[1]}"
        return inst_id
    
    def check_heartbeat(self):
        """检查心跳，如果超时则重连"""
        if not self.ws_connected:
            return
        
        elapsed = (datetime.now() - self.last_heartbeat).total_seconds()
        if elapsed > self.heartbeat_interval * 2:
            logger.warning(f"心跳超时 ({elapsed}秒)，尝试重连...")
            self._stop_websocket()
            self._start_websocket()
    
    def close(self):
        """关闭管理器"""
        logger.info("关闭行情数据管理器")
        self._stop_websocket()
        self.subscriptions.clear()
        self.latest_prices.clear()
        self.kline_cache.clear()
